import wfdb
import numpy as np
import pandas as pd
from scipy.signal import correlate, find_peaks, butter, filtfilt
from scipy.ndimage import uniform_filter1d

def bandpass_filter(signal, fs, lowcut, highcut, order=4):
    nyq = 0.5 * fs
    b, a = butter(order, [lowcut / nyq, highcut / nyq], btype='band')
    return filtfilt(b, a, signal)

def smooth_signal(signal, window_size=5):
    return uniform_filter1d(signal, size=window_size)

def remove_artifacts(signal, threshold=5.0):
    z = (signal - np.mean(signal)) / np.std(signal)
    artifact_indices = np.where(np.abs(z) > threshold)[0]
    clean_signal = signal.copy()
    for idx in artifact_indices:
        left = max(0, idx - 5)
        right = min(len(signal), idx + 6)
        clean_signal[idx] = np.median(signal[left:right])
    return clean_signal

def load_and_save_mimic_to_csv(record_name, output_csv='processed_data.csv'):
    record = wfdb.rdrecord(record_name)
    
    PPG_NAME = 'PLETH'
    ECG_NAME = 'II'
    ABP_NAME = 'ABP'
    
    assert PPG_NAME in record.sig_name, f"Required signal '{PPG_NAME}' not found. Available: {record.sig_name}"
    assert ABP_NAME in record.sig_name, f"Required signal '{ABP_NAME}' not found. Available: {record.sig_name}"
    assert ECG_NAME in record.sig_name, f"Required signal '{ECG_NAME}' not found. Available: {record.sig_name}"

    fs = record.fs
    signals = {
        'ppg': record.p_signal[:, record.sig_name.index(PPG_NAME)],
        'abp': record.p_signal[:, record.sig_name.index(ABP_NAME)],
        'ecg': record.p_signal[:, record.sig_name.index(ECG_NAME)]
    }

    for key in signals:
        sig = signals[key]
        if np.isnan(sig).any() or np.isinf(sig).any():
            print(f"信号 {key} 含有 NaN 或 Inf，正在进行修复...")
            mean_val = np.nanmean(sig[np.isfinite(sig)])
            signals[key] = np.nan_to_num(sig, nan=mean_val, posinf=mean_val, neginf=mean_val)

    # 信号处理（带通滤波 + 平滑 + 伪影剔除）
    signals['ppg'] = remove_artifacts(smooth_signal(bandpass_filter(signals['ppg'], fs, 0.5, 8)))
    signals['ecg'] = remove_artifacts(smooth_signal(bandpass_filter(signals['ecg'], fs, 0.5, 40)))
    signals['abp'] = remove_artifacts(smooth_signal(bandpass_filter(signals['abp'], fs, 0.5, 20)))

    time = np.arange(len(signals['ppg'])) / fs

    # PPG-ABP 对齐
    cross_corr = correlate(signals['ppg'] - signals['ppg'].mean(),
                           signals['abp'] - signals['abp'].mean(),
                           mode='full')
    delay = cross_corr.argmax() - (len(signals['ppg']) - 1)
    abp_aligned = np.roll(signals['abp'], -delay)

    # R峰检测
    ecg_peaks, _ = find_peaks(signals['ecg'], 
                              height=np.nanpercentile(signals['ecg'], 80),
                              distance=fs * 0.6)
    delta = np.mean(np.diff(ecg_peaks)) / fs if len(ecg_peaks) > 1 else 0.8

    # 构建DataFrame
    df = pd.DataFrame({
        'time': time,
        'pleth': signals['ppg'],
        'abp_raw': signals['abp'],
        'abp_aligned': abp_aligned,
        'ecg_ii': signals['ecg'],
        'r_peak': np.isin(np.arange(len(time)), ecg_peaks)
    })

    meta = (f"# fs={fs}, delay_samples={delay}, "
            f"ppg_name={PPG_NAME}, ecg_name={ECG_NAME}, "
            f"hr={60/delta:.1f}bpm\n")
    with open(output_csv, 'w') as f:
        f.write(meta)
        df.to_csv(f, index=False)

    print(f"数据已保存至 {output_csv}")
    print(f"信号长度: {len(signals['ppg'])/fs:.1f}秒 | 心率: {60/delta:.1f}bpm")

load_and_save_mimic_to_csv('3300784_0011', 'processed_data_0011.csv')